package com.liu.cloud.filter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;

import java.net.URI;

/**
 * 自定义负载均衡全局过滤器
 *
 * @author LIU
 * @version 1.0.0
 */
@Slf4j
//@Component
public class RequestHeaderLoadBalancerClientFilter extends LoadBalancerClientFilter {

    private static final String TERMINAL_ID = "TerminalID";

    public RequestHeaderLoadBalancerClientFilter(LoadBalancerClient loadBalancer, LoadBalancerProperties properties) {
        super(loadBalancer, properties);
    }

    @Override
    protected ServiceInstance choose(ServerWebExchange exchange) {
        // 获取请求中的cookie
        String wellId = exchange.getRequest().getHeaders().getFirst(TERMINAL_ID);
        if (StringUtils.isEmpty(wellId)) {
            return super.choose(exchange);
        }
        if (this.loadBalancer instanceof RibbonLoadBalancerClient) {
            RibbonLoadBalancerClient client = (RibbonLoadBalancerClient) this.loadBalancer;
            Object attrValue = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
            Assert.notNull(attrValue,"gatewayRequestUrl值为空！");
            String serviceId = ((URI) attrValue).getHost();

            // 这里使用请求头做为选择服务实例的key
            ServiceInstance choose = client.choose(serviceId, wellId);
//            log.info("wellId:{},serviceId:{},choose:{}",wellId,serviceId,choose);
            return choose;
        }
        return super.choose(exchange);
    }
}
